from vec import Vec2
import math
import random
import pgzrun

WIDTH = 750
HEIGHT = 750

PEN_COLOR = (0, 200, 0)
BACK_COLOR = (0, 0, 0)
ptList = []
length = 650
pt_mid = Vec2(WIDTH / 2, HEIGHT / 2)
y_offset=length/math.sqrt(3) 
pt_a=Vec2(pt_mid.x,pt_mid.y-y_offset)
pt_d=Vec2(pt_mid.x,pt_mid.y+length/2/math.sqrt(3))
pt_b=Vec2(pt_d.x-length/2,pt_d.y)
pt_c=Vec2(pt_d.x+length/2,pt_d.y)
ptList.append(pt_a)
ptList.append(pt_b)
ptList.append(pt_c)

def distort(p):
    r = random.randint(1, 3)
    vx = random.randint(-100, 100)
    vy = random.randint(-100, 100)
    v = Vec2(vx, vy)
    v.normalize()
    dp = p + v * r

    return dp

def draw_fractal(pa, pb,m):
    if m == 1:
        if enableDistort == False:
            screen.draw.line(pa.raw_get(), pb.raw_get(), PEN_COLOR)
        else:
            screen.draw.line(distort(pa).raw_get(), distort(pb).raw_get(), PEN_COLOR)
    else:
        pts=[]
        pt1=pa*2/3+pb/3
        midPt=(pa+pb)/2
        dir=pb-pa
        long=dir.length()
        dir.normalize()
        vdir=Vec2(0,0)
        vdir.x=-dir.y
        vdir.y=dir.x
        dist=long/3/2*math.sqrt(3)
        pt2=midPt+vdir*dist
        pt3=pa/3+pb*2/3
        pts.append(pa)
        pts.append(pt1)
        pts.append(pt2)
        pts.append(pt3)
        pts.append(pb)
        for i in range(len(pts) - 1):
            draw_fractal(pts[i], pts[i + 1], m - 1)

count = len(ptList)
n = 7
enableDistort = False

def on_key_down(key):
    global enableDistort
    if key == keys.ESCAPE:
        quit() 
    elif key == keys.SPACE:
        enableDistort = not enableDistort

def update():
    pass

def draw():
    screen.fill(BACK_COLOR)
    
    for i in range(count):
        if i == count - 1:
            draw_fractal(ptList[count - 1], ptList[0],n)
        else:
            draw_fractal(ptList[i], ptList[i + 1],n) 

pgzrun.go()
